#include"include.h"

int sendNetMessage(peer_t *p, netmsg_t msg)
{
   if(!p->ready)
      return 1;
   loggerLog(LOG_INFO, "\tSending network message.\n");
   int ret = 0;

   int pllen = msg.len_prefix - 1;
   void *sendBuf = NULL;
   if(pllen < 0)
      pllen = 0;

   // Allocate the buffer being sent
   int msglen = 5 + pllen;
   sendBuf = malloc(msglen);
   // Properly fill message buffer
   netmsg_t *ptr = sendBuf;
   ptr->len_prefix = htonl(msg.len_prefix);
   ptr->message_ID = msg.message_ID;
   memcpy(sendBuf+5, msg.payload, pllen);




   // Send message buffer
   if(send(p->socket_descriptor, sendBuf, msglen, 0) != msglen)
   {
      if(errno == EPIPE)
      {
         ret = PEER_SHUTDOWN;
         loggerLog(LOG_INFO, "\tThe peer has disconnected.");
      }
      else
         ret = SOCK_ERR; // Error sending message
   }

   // Free used buffer
   free(sendBuf);
   
   if(ret)
      loggerLog(LOG_ERR, "\tError sending network message.\n");
   else
      loggerLog(LOG_INFO, "\tNetwork message sent successfully.\n");

   return ret;
}

int peerHandshake(peer_t *p, char *peer_id, unsigned char info_hash[20])
{
   loggerLog(LOG_INFO, "Connecting to peer.\n");
   if((p->socket_descriptor = socket(PF_INET, SOCK_STREAM, IPPROTO_TCP)) < 0)
   {
      loggerLog(LOG_ERR, "Could not create socket.\n");
      return 1;
   }
   struct sockaddr_in peerAddr;
   memset(&peerAddr, 0, sizeof(peerAddr));
   peerAddr.sin_family = AF_INET;
   peerAddr.sin_addr.s_addr = htonl(p->ip);
   peerAddr.sin_port = htons(p->port);

   if(connect(p->socket_descriptor, (struct sockaddr *) &peerAddr, sizeof(peerAddr)) < 0)
   {
      loggerLog(LOG_ERR, "Could not connect to peer %x.\n", p->ip);
      return 1;
   }

   int sock = p->socket_descriptor;
   // Send a handshake
   int err = sendHandshake(sock, peer_id, info_hash);
   // Now receive one
   if(!err)
   {
      p->my_alive = gettime();
      loggerLog(LOG_INFO, "Receiving handshake from peer.\n");
      unsigned char pinfo_hash[20];

      // Get message length
      uint8_t pstrlen;
      recvData(sock, &pstrlen, 1);

      // Get pstr
      char *useless = malloc(pstrlen);
      recvData(sock, useless, pstrlen);
      free(useless);

      // Receive the reserved bytes
      uint8_t reserved[8];
      recvData(sock, &reserved, 8);
      //Ignored for now; we don't deal with extensions

      // Receive the infohash
      recvData(sock, pinfo_hash, 20);

      // Receive peer_id
      char ppeer_id[20];
      recvData(sock, ppeer_id, 20);

      // Add peer_id to peer_t
      memcpy(p->peer_id, ppeer_id, 20);

      int logi;
      loggerLog(LOG_INFO, "Handshake received from peer ");
      for(logi = 0; logi < 20; logi++)
         loggerLog(LOG_INFO, "%x", peer_id[logi]);
      loggerLog(LOG_INFO, ".\n");
      p->peer_alive = gettime();
   }
   p->ready = 1;

   return err;
}

int sendHandshake(int sock, char *peer_id, unsigned char info_hash[20])
{

   loggerLog(LOG_INFO, "Sending handshake to peer.\n");
   // Success indicator
   int ret = 0;

   // Assume BitTorrent Protocol version 1.0
   unsigned char pstrlen = 19;         // Assumed for version 1.0
   char *pstr = "BitTorrent protocol"; // Assumed for version 1.0
   uint64_t reserved = 0;              // All current impls use 0s

   int msgbuffsize = 49+pstrlen;

   // Construct buffer to send
   unsigned char msg[msgbuffsize];
   void *loc = msg+1;

   // Fill buffer to send
   msg[0] = pstrlen;               // Fill <pstrlen> field
   memcpy(loc, pstr, pstrlen);     // Fill <pstr> field
   loc += pstrlen;
   *((uint64_t *) loc) = reserved; // Fill <reserved> field
   loc += 8;
   memcpy(loc, info_hash, 20);     // Fill <info_hash> field
   loc += 20;
   memcpy(loc, peer_id, 20);       // Fill <peer_id> field

   // Send message buffer
   if(send(sock, msg, msgbuffsize, 0) != msgbuffsize)
      ret = 1; // Error sending message

   if(ret)
      loggerLog(LOG_ERR, "Error sending handshake!\n");
   else
      loggerLog(LOG_INFO, "Handshake successfully sent\n");
   
   return ret;
}

int peerSendKeepAlive(peer_t *p)
{
   loggerLog(LOG_INFO, "Sending Keep-Alive.\n");
   // Construct proper netmsg_t
   netmsg_t toSend = {0, KEEPALIVE, NULL};
   return sendNetMessage(p, toSend);
}

int peerSendChoke(peer_t *p)
{
   loggerLog(LOG_INFO, "Sending Choke.\n");
   // Construct proper netmsg_t
   netmsg_t toSend = {1, CHOKE, NULL};
   return sendNetMessage(p, toSend);
}

int peerSendUnchoke(peer_t *p)
{
   loggerLog(LOG_INFO, "Sending Unchoke.\n");
   // Construct proper netmsg_t
   netmsg_t toSend = {1, UNCHOKE, NULL};
   return sendNetMessage(p, toSend);
}

int peerSendInterested(peer_t *p)
{
   loggerLog(LOG_INFO, "Sending Interested.\n");
   // Construct proper netmsg_t
   netmsg_t toSend = {1, INTERESTED, NULL};
   return sendNetMessage(p, toSend);
}

int peerSendNotInterested(peer_t *p)
{
   loggerLog(LOG_INFO, "Sending Not Interested.\n");
   // Construct proper netmsg_t
   netmsg_t toSend = {1, NOTINTERESTED, NULL};
   return sendNetMessage(p, toSend);
}

int peerSendHave(peer_t *p, uint32_t bid)
{  
   loggerLog(LOG_INFO, "Sending Have.\n");
   // Construct proper netmsg_t
   bid = htonl(bid);
   netmsg_t toSend = {5, HAVE, &bid};
   return sendNetMessage(p, toSend);
}

int peerSendBitfield(peer_t *p, bitfield_t *bitfield)
{
   loggerLog(LOG_INFO, "Sending Bitfield.\n");
   // Construct proper netmsg_t
   netmsg_t toSend = {1+bitfield->length, BITFIELD, bitfield->bitfield};
   return sendNetMessage(p, toSend);
}

int peerSendRequest(peer_t *p, uint32_t ndx, uint32_t begin,
   uint32_t len)
{
   loggerLog(LOG_INFO, "Sending Request.\n");
   // Construct proper netmsg_t
   uint32_t flds[] = {htonl(ndx), htonl(begin), htonl(len)};
   netmsg_t toSend = {13, REQUEST, flds};
   return sendNetMessage(p, toSend);
}

int peerSendPiece(peer_t *p, uint32_t ndx, uint32_t begin,
   unsigned char *block, int block_len)
{
   loggerLog(LOG_INFO, "Sending Piece with index %d.\n", ndx);
   int ret = 1;
   // Construct proper netmsg_t
   // Using heap due to potentially large blocks
   void *buf = malloc(8+block_len);
   // Construct message buffer
   *((uint32_t *) buf) = htonl(ndx);       // <index> field
   *((uint32_t *) (buf+4)) = htonl(begin); // <begin> field
   memcpy(buf+8, block, block_len);        // <block> field

   netmsg_t toSend = {9+block_len, PIECE, buf};

   ret = sendNetMessage(p, toSend); // Send the message
   free(buf);
   return ret;
}

int peerSendCancel(peer_t *p, uint32_t ndx, uint32_t begin,
   uint32_t len)
{
   loggerLog(LOG_INFO, "Sending Cancel.\n");
   // Construct proper netmsg_t
   uint32_t flds[] = {htonl(ndx), htonl(begin), htonl(len)};
   netmsg_t toSend = {13, CANCEL, flds};
   return sendNetMessage(p, toSend);
}

int peerSendPort(peer_t *p, uint16_t listen_port)
{
   loggerLog(LOG_INFO, "Sending Port.\n");
   // Construct proper netmsg_t
   listen_port = htons(listen_port);
   netmsg_t toSend = {3, PORT, &listen_port};
   return sendNetMessage(p, toSend);
}

int peerCheckRecieved(peer_t *p, void *torrent, void (*callback) (void *torrent, peer_t *, netmsg_t))
{
   loggerLog(LOG_INFO, "Checking for received data...\n");
   netmsg_t msg;
   int left;
   struct pollfd sockPoll = {p->socket_descriptor, POLLIN | POLLPRI, 0};

   // Check for data
   int res;
   if((res = poll(&sockPoll, 1, CHECK_WAIT_TIME)) <= 0)
   {
      loggerLog(LOG_INFO, "\tNo network data here.\n");
      return 0;
   }
   else
   {  // There's data to read
      // Read length of message
      loggerLog(LOG_INFO, "\tNetwork data waiting.\n");
      int tmp;
      if((tmp = recv(p->socket_descriptor, &(msg.len_prefix), 4, 0)) != 4)
      {
         loggerLog(LOG_ERR, "Peer disconnected on socket %d.\n",
            p->socket_descriptor);
         return PEER_SHUTDOWN; // Error receiving data from socket
      }
      msg.len_prefix = ntohl(msg.len_prefix); // Decode length
      left = msg.len_prefix - 1;
      // Read message opcode
      if(msg.len_prefix != 0)
      {
         if(recv(p->socket_descriptor, &(msg.message_ID), 1, 0) != 1)
         {
            loggerLog(LOG_ERR, "Socket error on recv id\n");
            return SOCK_ERR; // Error receiving data from socket
         }
      }
      else
         msg.message_ID = KEEPALIVE;

      if(left > 0)
      {  // Check to see if there is a payload
         msg.payload = malloc(left);
         // Read payload
         if(recvData(p->socket_descriptor, msg.payload, left))
         {
            loggerLog(LOG_ERR, "Socket error on recv payload\n");
            free(msg.payload);
            return SOCK_ERR; // Error receiving data from socket
         }
      }
      else
         msg.payload = NULL;

      // Call the callback func
      loggerLog(LOG_INFO, "\tCall function with data.\n");
      callback(torrent, p, msg);

      if(poll(&sockPoll, 1, CHECK_WAIT_TIME) <= 0)
         return 0;
      else
         return NOT_EMPTY;
   }
}

peer_t *peerInit(unsigned long torrent_len, unsigned int pieceLength)
{
   peer_t *p = malloc(sizeof(peer_t));
   p->my_interest = 0;
   p->peer_interest = 0;
   p->my_choking = 1;
   p->peer_choking = 1;
   p->last_choke = 0;
   p->peer_alive = gettime();
   p->my_alive = gettime();
   p->my_reqs = 0;
   p->ready = 0;

   p->requestList = createList();
   p->bitfield = bitfieldInit(torrent_len);
   p->activeReqs = newDict((int (*)(const void *, const void *))integerCompare);
   p->pieceLength = pieceLength;

   return p;
}
void peerDisconnect(peer_t *p)
{
   p->ready = 0;
   p->my_reqs = 0;
   p->peer_choking = 1;
   p->my_choking = 1;
   p->my_interest = 0;
   p->peer_interest = 0;
}
void peerPrint(peer_t *p)
{
   struct in_addr toprint;
   toprint.s_addr = htonl(p->ip);
   printf("Peer %s:%d\n", inet_ntoa(toprint), p->port);
   printf("my interest / their interest / my choke / their choke = %d / %d / %d / %d\n", p->my_interest, p->peer_interest, p->my_choking, p->peer_choking);
}

void peerSetHave(peer_t *p, unsigned int ndx)
{
   bitfieldSet(p->bitfield, ndx, 1);
}
int peerGetHave(peer_t *p, unsigned int ndx)
{
   return bitfieldGet(p->bitfield, ndx);
}
void peerSetBitfield(peer_t *p, unsigned int len, unsigned char *field)
{
   int i;
   unsigned int numBits = bitfieldGetLen(p->bitfield);
   for(i=0;i<numBits;i++)
   {
      bitfieldSet(p->bitfield, i, field[i/8] |= ((1 << (i%8))));
   }
}

int peerRequested(peer_t *p, unsigned int ndx, unsigned int offset)
{
   if(!dictContainsKey(p->activeReqs, &ndx))
      return 0;
   bitfield_t *b = dictGet(p->activeReqs, &ndx);
   return bitfieldGet(b, offset / BLOCK_SIZE);
}

void peerSetRequest(peer_t *p, unsigned int ndx, unsigned int offset, unsigned int len)
{
   p->my_reqs++;
   bitfield_t *b;
   if(dictContainsKey(p->activeReqs, &ndx))
   {
      b = dictGet(p->activeReqs, &ndx);
      bitfieldSet(b, offset / BLOCK_SIZE, 1);
   }
   else
   {
      b = bitfieldInit(p->pieceLength / BLOCK_SIZE);
      bitfieldSet(b, offset / BLOCK_SIZE, 1);
      dictPut(p->activeReqs, &ndx, b);
   }
   /*piece_request *rq = malloc(sizeof(piece_request));
   rq->ndx = ndx;
   rq->offset = offset;
   rq->len = len;
   listLock(p->requestList);
   addObject(p->requestList, rq);
   listRelease(p->requestList);*/
}
void peerCancelRequest(peer_t *p, unsigned int ndx, unsigned int offset, unsigned int len)
{
   piece_request *rq;
   listLock(p->requestList);
   for(resetList(p->requestList);!endList(p->requestList);)
   {
      rq = getObject(p->requestList);
      if(rq->len == len && rq->ndx == ndx && rq->offset == offset)
      {
         removeObject(p->requestList, rq);
         break;
      }
   }
   listRelease(p->requestList);
}

